package edu.berkeley.nlp.assignments;

import java.util.List;
import java.util.Map;

import edu.berkeley.nlp.classify.LabeledInstance;
import edu.berkeley.nlp.classify.ProbabilisticClassifier;
import edu.berkeley.nlp.classify.ProbabilisticClassifierFactory;
import edu.berkeley.nlp.util.Counter;
import edu.berkeley.nlp.util.CounterMap;
import edu.berkeley.nlp.util.MapFactory;

public class ClassConditionalLMClassifier<I, L> implements
		ProbabilisticClassifier<I, L> {

	public static class Factory<I, L> implements
			ProbabilisticClassifierFactory<I, L> {

		public ProbabilisticClassifier<I, L> trainClassifier(
				List<LabeledInstance<I, L>> trainingData, int type, boolean usingAlpha) {

			CounterMap<L, Character> classCounter = new CounterMap<L, Character>();
			
			if (type == WORD_UNIGRAM) {
				CounterMap<L, String> wordCounter = new CounterMap<L, String>();
				for (LabeledInstance<I, L> datum : trainingData) {
					L classLabel = datum.getLabel();
					
					String word = (String)datum.getInput();
					
					wordCounter.incrementCount(classLabel, word, 1.0);
					
				}
				
				return new ClassConditionalLMClassifier<I, L>(wordCounter, type, 0);
			} 
			
			if (type == CHARACTER_UNIGRAM || type == CHARACTER_BIGRAM) {
				
				for (LabeledInstance<I, L> datum : trainingData) {
					L classLabel = datum.getLabel();
					
					String word = (String)datum.getInput();
					
					for (int i = 0; i < word.length(); i ++) {
						classCounter.incrementCount(classLabel, word.charAt(i), 1.0);
					}
				}
			} 
			
			
			if (type == CHARACTER_BIGRAM) {
				Map<L, CounterMap<Character, Character>> bigramClassMap = new MapFactory.HashMapFactory<L, CounterMap<Character, Character>>().buildMap();

				for (LabeledInstance<I, L> datum : trainingData) {
					L classLabel = datum.getLabel();
					
					String word = (String)datum.getInput();
					if (!bigramClassMap.containsKey(classLabel))
						bigramClassMap.put(classLabel, new CounterMap<Character, Character>());
					
					for (int i = 1; i < word.length(); i ++) {
						bigramClassMap.get(classLabel).incrementCount(word.charAt(i-1), word.charAt(i), 1.0);
					}
					
				}
				return new ClassConditionalLMClassifier<I, L>(classCounter, bigramClassMap, type, usingAlpha);
			}

			if (type == CHARACTER_UNIGRAM) {
				return new ClassConditionalLMClassifier<I, L>(classCounter, type);
				
			} 
				
			return null;
		}

		public ProbabilisticClassifier<I, L> trainClassifier(
				List<LabeledInstance<I, L>> trainingData) {
			// TODO Auto-generated method stub
			// obsolete!!!
			return null;
		}

	}

	public static int CHARACTER_UNIGRAM = 0;
	public static int CHARACTER_BIGRAM = 1;
	public static int WORD_UNIGRAM = 2;
	int type;

	Counter<String> unkCount = new Counter<String>();
	Counter<String> totalCount = new Counter<String>();
	
	//character unigram
	CounterMap<L, Character>  classCounter;
	//character bigram
	Map<L, CounterMap<Character, Character>> bigramClassMap;
	//word unigram
	CounterMap<L, String>  wordCounter;
	boolean usingAlpha;
	double alpha;
	public void setAlpha(double alpha) {
		this.alpha = alpha;
	}

	public ClassConditionalLMClassifier(CounterMap<L, String>  wordCounter, int type, int dumb) {
		this.type = type;
		this.wordCounter = wordCounter;
	}
	
	public ClassConditionalLMClassifier(CounterMap<L, Character>  classCounter, int type) {
		this.type = type;
		this.classCounter = classCounter;
	}
	
	public ClassConditionalLMClassifier(CounterMap<L, Character>  classCounter, Map<L, CounterMap<Character, Character>> bigramClassMap, int type, boolean usingAlpha) {
		this.type = type;
		this.bigramClassMap = bigramClassMap;
		this.classCounter = classCounter;
		this.usingAlpha = usingAlpha;
	}

	public Counter<L> getProbabilities(I instance) {
	    Counter<L> counter = new Counter<L>();
		if (type == CHARACTER_UNIGRAM) {
			String word = (String) instance;
			for (L label: classCounter.keySet()) {
				double prob = 1.0;
				for (int i = 0; i < word.length(); i++) {
					prob *= getProb(label, word.charAt(i));
				}
				counter.incrementCount(label, prob);
			}
		    counter.normalize();
			return counter;
		} else if (type == CHARACTER_BIGRAM) {
			String word = (String) instance;
			for (L label: classCounter.keySet()) {
				double prob = 1.0;
				prob *= getProb(label, word.charAt(0));
				for (int i = 1; i < word.length(); i++) {
					prob *= getProb(label, word.charAt(i-1), word.charAt(i));
				}
				counter.incrementCount(label, prob);
			}
		    counter.normalize();
			return counter;
			
		} else if (type == WORD_UNIGRAM) {
			for (L label: wordCounter.keySet()) {
				double prob = 0, count, total = wordCounter.getCounter(label).totalCount();
				if ((count = wordCounter.getCount(label, (String)instance)) == 0) {
//					unkCount.incrementCount((String)correctLabel, 1.0);
					prob = 100.0 / (total + 1.0);
				} else
					prob = count / (total + 1.0);
				
//				totalCount.incrementCount((String)correctLabel, 1.0);
				
				counter.incrementCount(label, prob);
			}
			counter.normalize();
			return counter;
		}
		return null;
	}

	public L getLabel(I instance) {
	    Counter<L> counter = getProbabilities(instance);
	    return counter.argMax();
	    
	}
	
	
	// CHARACTER_UNIGRAM
	public double getProb(L label, Character c) {
		double total = classCounter.getCounter(label).totalCount();
		if (classCounter.getCount(label, c) == 0) 
			return 1.0 / (total + 1.0);
		else
			return classCounter.getCount(label, c) / (total + 1.0);
	}
	
	
	// CHARACTER_BIGRAM
	public double getProb(L label, Character lastC, Character c) {

		double total = bigramClassMap.get(label).getCounter(lastC).totalCount();
		if (usingAlpha) {
			double prob = 0.0;
			prob += alpha * bigramClassMap.get(label).getCount(lastC, c)
			/ (total + 1.0);
			prob += (1.0 - alpha) * getProb(label, c);
			return prob;
		} else {
			if (bigramClassMap.get(label).getCount(lastC, c) == 0)
				// / backoff to unigram
				return getProb(label, c);
			// / additive smoothing
			// return 1.0 / (total + 1.0);
			else
				return bigramClassMap.get(label).getCount(lastC, c)
						/ (total + 1.0);
		}
	}
	
	
	
	public void printUNKStats(List<LabeledInstance<String, String>> testData) {
		for (LabeledInstance<String, String> testDatum : testData) {
		      String name = testDatum.getInput();
		      String correctLabel = testDatum.getLabel();
//		      for (L label: wordCounter.keySet()) {
		    	  if (wordCounter.getCount((L)correctLabel, name) == 0) {
		    		  unkCount.incrementCount((String)correctLabel, 1.0);
//		    		  break;
		    	  }
//		      }
	    	  totalCount.incrementCount((String)correctLabel, 1.0);
		}
		for (L label: wordCounter.keySet()) {
			System.out.println((String)label +  unkCount.getCount((String)label) / totalCount.getCount((String)label));
		}
	}

}
